package com.yangzhe.algorithm.c018;

import java.util.ArrayList;
import java.util.List;
import java.util.Stack;

// 不用递归，用迭代的方式实现二叉树的三序遍历
public class BinaryTreeTraversalIteration {

    public static class TreeNode {
        public int val;
        public TreeNode left;
        public TreeNode right;

        public TreeNode(int v) {
            val = v;
        }
    }

    /**
     * 先序遍历
     *
     * @param head
     */
    public static void preOrder(TreeNode head) {
        Stack<TreeNode> stack = new Stack<>();
        stack.push(head);
        while (!stack.isEmpty()) {
            TreeNode pop = stack.pop();
            System.out.print(pop.val + ",");

            if (pop.right != null) {
                stack.push(pop.right);
            }

            if (pop.left != null) {
                stack.push(pop.left);
            }
        }
    }

    /**
     * 中序
     *
     * @param head
     */
    public static void inOrder(TreeNode head) {
        Stack<TreeNode> stack = new Stack<>();
        while (!stack.isEmpty() || head != null) {
            if (head != null) {
                // 将节点左树全都放到栈中
                stack.push(head);
                head = head.left;
            } else {
                // 弹出打印
                head = stack.pop();
                System.out.print(head.val + ",");

                // 右节点子树继续处理左节点
                head = head.right;
            }
        }
    }

    /**
     * 后序
     *
     * @param head
     */
    public static void posOrderTwoStacks(TreeNode head) {
        Stack<TreeNode> stack = new Stack<>();
        Stack<Integer> result = new Stack<>();
        stack.push(head);
        while (!stack.isEmpty()) {
            // 先序是 中左右， 这里处理成 中右左，入栈再出栈变成后序 左右中
            TreeNode pop = stack.pop();
            result.push(pop.val);
            if (pop.left != null) {
                stack.push(pop.left);
            }

            if (pop.right != null) {
                stack.push(pop.right);
            }
        }

        while (!result.isEmpty()) {
            System.out.print(result.pop() + ",");
        }
    }

    public static void posOrderOneStack(TreeNode head) {
        Stack<TreeNode> stack = new Stack<>();
        stack.push(head);
        if (head != null) {
            while (!stack.isEmpty()) {
                TreeNode current = stack.peek();
                if (current.left != null && current.left != head && current.right != head) {
                    stack.push(current.left);
                } else if (current.right != null && current.right != head) {
                    stack.push(current.right);
                } else {
                    head = stack.pop();
                    System.out.print(head.val + ",");
                }
            }
        }
    }

    public static void main(String[] args) {
        TreeNode head = new TreeNode(1);
        head.left = new TreeNode(2);
        head.right = new TreeNode(3);
        head.left.left = new TreeNode(4);
        head.left.right = new TreeNode(5);
        head.right.left = new TreeNode(6);
        head.right.right = new TreeNode(7);
        preOrder(head);
        System.out.println("先序遍历非递归版");
        inOrder(head);
        System.out.println("中序遍历非递归版");
        posOrderTwoStacks(head);
        System.out.println("后序遍历非递归版 - 2个栈实现");
        posOrderOneStack(head);
        System.out.println("后序遍历非递归版 - 1个栈实现");
    }

    // 用一个栈完成先序遍历
    // 测试链接 : https://leetcode.cn/problems/binary-tree-preorder-traversal/
    public static List<Integer> preorderTraversal(TreeNode head) {
        List<Integer> ans = new ArrayList<>();
        Stack<TreeNode> stack = new Stack<>();
        stack.push(head);
        while (head != null && !stack.isEmpty()) {
            TreeNode cur = stack.pop();
            ans.add(cur.val);
            if (cur.right != null) {
                stack.push(cur.right);
            }

            if (cur.left != null) {
                stack.push(cur.left);
            }
        }
        return ans;
    }

    // 用一个栈完成中序遍历
    // 测试链接 : https://leetcode.cn/problems/binary-tree-inorder-traversal/
    public static List<Integer> inorderTraversal(TreeNode head) {
        List<Integer> ans = new ArrayList<>();
        Stack<TreeNode> stack = new Stack<>();
        while (head != null || !stack.isEmpty()) {
            if (head != null) {
                stack.push(head);
                head = head.left;
            } else {
                TreeNode pop = stack.pop();
                ans.add(pop.val);
                head = pop.right;
            }
        }

        return ans;
    }

    // 用两个栈完成后序遍历
    // 提交时函数名改为postorderTraversal
    // 测试链接 : https://leetcode.cn/problems/binary-tree-postorder-traversal/
    public static List<Integer> postorderTraversalTwoStacks(TreeNode root) {
        List<Integer> ans = new ArrayList<>();
        Stack<TreeNode> stack = new Stack<>();
        Stack<Integer> result = new Stack<>();
        if (root == null) {
            return ans;
        }

        stack.push(root);
        while (!stack.isEmpty()) {
            TreeNode cur = stack.pop();
            result.push(cur.val);
            if (cur.left != null) {
                stack.push(cur.left);
            }

            if (cur.right != null) {
                stack.push(cur.right);
            }
        }
        while(!result.isEmpty()) {
            ans.add(result.pop());
        }

        return ans;
    }

    // 用一个栈完成后序遍历
    // 提交时函数名改为postorderTraversal
    // 测试链接 : https://leetcode.cn/problems/binary-tree-postorder-traversal/
    public static List<Integer> postorderTraversalOneStack(TreeNode root) {
        List<Integer> ans = new ArrayList<>();
        Stack<TreeNode> stack = new Stack<>();
        if (root == null) {
            return ans;
        }

        TreeNode head = root;
        stack.push(head);
        while(!stack.isEmpty()) {
            TreeNode cur = stack.peek();
            if (cur.left != null && cur.left != head && cur.right != head) {
                stack.push(cur.left);
            } else if (cur.right != null && cur.right != head) {
                stack.push(cur.right);
            } else {
                head = stack.pop();
                ans.add(head.val);
            }
        }

        return ans;
    }

}
